package sketch.instrumenter;

import java.util.Arrays;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import org.eclipse.jdt.core.dom.AST;
import org.eclipse.jdt.core.dom.ASTNode;
import org.eclipse.jdt.core.dom.AbstractTypeDeclaration;
import org.eclipse.jdt.core.dom.Block;
import org.eclipse.jdt.core.dom.CatchClause;
import org.eclipse.jdt.core.dom.CompilationUnit;
import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.ExpressionStatement;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.MethodInvocation;
import org.eclipse.jdt.core.dom.SimpleName;
import org.eclipse.jdt.core.dom.SingleVariableDeclaration;
import org.eclipse.jdt.core.dom.Statement;
import org.eclipse.jdt.core.dom.StringLiteral;
import org.eclipse.jdt.core.dom.TryStatement;
import org.eclipse.jdt.internal.core.dom.NaiveASTFlattener;

import sketch.ast.ASTUtils;
import sketch.main.SketchTestOptions;
import sketch.ounit.Observer;
import sketch.specs.SketchedTryCatchBlock;
import sketch.util.Checker;

public class SourceCodeInstrumentingVisitor extends NaiveASTFlattener {
	
	public final CompilationUnit unit;
	
	public final Collection<TryStatement> instrumentedTries;
	public final Collection<MethodInvocation> instrumentedObservers;
	
	public static String TRACER_CLASS = "sketch.instrumenter.Tracer";
	public static String START_METHOD = "startMethod";
	public static String END_METHOD = "endMethod";
	public static String RECORD_FAIL_POS_METHOD = "recordFailPosition";
	public static String RECORD_EXCEPTION_METHOD = "recordException";
	public static String RECORD_OBSERVATION = "recordObservation";
	
	public static String SEP_SYMBOL = "#";
	
	public SourceCodeInstrumentingVisitor(CompilationUnit unit) {
		Checker.checkNull(unit, "The compilation unit could not be null!");
		this.unit = unit;
		this.instrumentedTries = new LinkedList<TryStatement>();
		this.instrumentedObservers = new LinkedList<MethodInvocation>();
	}
	
	public boolean hasInstrumented() {
		if(this.instrumentedTries.isEmpty() && this.instrumentedObservers.isEmpty()) {
			return false;
		}
		return true;
	}
	
//	@Override
//	public boolean visit(CompilationUnit unit) {
//		//do some thing here
//		boolean result = super.visit(unit);
//		return result;
//	}
	
	@Override
	public boolean visit(TryStatement tryStatement) {
		if(isWrappedTryCatch(tryStatement)) {
			//int startPosition = tryStatement.getStartPosition();
			//do something here
			//add a new statement to the try block
			Statement tryInstrument = this.createInstrumentationStatement(tryStatement, TRACER_CLASS,
					RECORD_FAIL_POS_METHOD);
			tryStatement.getBody().statements().add(tryInstrument);
			
			//add a new statement to the catch block
			Checker.checkTrue(tryStatement.catchClauses().size() == 1, "The size of catchCaluse is not 1, it is: "
					+ tryStatement.catchClauses().size());
			CatchClause cc = (CatchClause)tryStatement.catchClauses().get(0);
			Statement catchInstrument =
				this.createInstrumentationStatement(tryStatement, TRACER_CLASS, RECORD_EXCEPTION_METHOD, SketchedTryCatchBlock.EXCEPTION_NAME);
			cc.getBody().statements().add(catchInstrument);
			
			//add to the collection
			this.instrumentedTries.add(tryStatement);
		}
		return super.visit(tryStatement);
	}
	
	private Statement createInstrumentationStatement(ASTNode nodeToTrace, String className, String staticMethodName,
			String...parameters)	{
		AST ast = nodeToTrace.getAST();
		MethodInvocation tracerInvoke = ast.newMethodInvocation();
		tracerInvoke.setExpression(ast.newName(className));
		tracerInvoke.setName(ast.newSimpleName(staticMethodName));
		List<Expression> arguments = tracerInvoke.arguments();
//		arguments.add(ast.newNumberLiteral( String.valueOf(this.unit.getLineNumber(nodeToTrace.getStartPosition()))));
		arguments.add(ast.newNumberLiteral( String.valueOf(nodeToTrace.getStartPosition())));
		for(String param : parameters) {
		    SimpleName name = ast.newSimpleName(param);
		    arguments.add(name);
		}
		Statement insertStatement = ast.newExpressionStatement(tracerInvoke);
		
		return insertStatement;
	}
	
	@Override
	public boolean visit(MethodDeclaration method) {
		if(!ASTUtils.hasOUnitAnnotation(method)) {
			return false;
		}
		
		Checker.checkTrue(unit.getAST() == method.getAST(), "The ast should be equal.");
		if(SketchTestOptions.verbose) {
		    System.out.println("instrumenting method: " + method.getName());
		}
		MethodInvocation startTracer = method.getAST().newMethodInvocation();
		startTracer.setExpression(method.getAST().newName(TRACER_CLASS));
		startTracer.setName(method.getAST().newSimpleName(START_METHOD));
		List<Expression> startArguments = startTracer.arguments();
		StringLiteral stringLiteral = method.getAST().newStringLiteral();
//		String methodStartString = method.getName().toString();
		String methodStartString = this.constructMethodSignature(method, this.unit);
		stringLiteral.setLiteralValue(methodStartString);
		startArguments.add(stringLiteral);
		
		Statement insertStartStatement = method.getAST().newExpressionStatement(startTracer);
		method.getBody().statements().add(0, insertStartStatement);
		
		//The end
		MethodInvocation endTracer = method.getAST().newMethodInvocation();
		endTracer.setExpression(method.getAST().newName(TRACER_CLASS));
		endTracer.setName(method.getAST().newSimpleName(END_METHOD));
		List<Expression> endArguments = endTracer.arguments();
		stringLiteral = method.getAST().newStringLiteral();
//		String methodEndString = method.getName().toString();
		String methodEndString = this.constructMethodSignature(method, this.unit);
		stringLiteral.setLiteralValue(methodEndString);
		endArguments.add(stringLiteral);
		
		Statement insertEndStatement = method.getAST().newExpressionStatement(endTracer);
		method.getBody().statements().add(insertEndStatement);
		
		return super.visit(method);
	}
	
	public static String constructMethodSignature(MethodDeclaration method, CompilationUnit unit) {
		String packageName = unit.getPackage() == null ? "" : unit.getPackage().getName().toString();
		AbstractTypeDeclaration class_type = ASTUtils.getParent(method, AbstractTypeDeclaration.class);
		String methodName = method.getName().toString();
		Checker.checkNull(class_type, "The class for method: " + methodName + " could not be null.");
		
		return packageName + SEP_SYMBOL + class_type.getName().toString() + SEP_SYMBOL + methodName; 
	}
	
	public static String[] parseIntoPackageClassMethodNames(String s) {
		Checker.checkNull(s, "The string to parse could not be null.");
		String[] names = s.split(SEP_SYMBOL);
		Checker.checkTrue(names.length == 3, "The string s is not valid: " + s);
		return names;
	}
	
	@Override
	public boolean visit(Block block) {
		Checker.checkTrue(unit.getAST() == block.getAST(), "The ast should be equal.");
		
		List<Statement> statements = block.statements();
		Map<Integer, Statement> insertingStatements = new LinkedHashMap<Integer, Statement>();
		int i = 0;
		for(Statement statement : statements) {
			if(statement instanceof ExpressionStatement) {
				ExpressionStatement exprStatement = (ExpressionStatement)statement;
				Expression expression = exprStatement.getExpression();
				//record the value of the expression
				if(expression instanceof MethodInvocation && isObserverExpression((MethodInvocation)expression)) {
					if(SketchTestOptions.verbose) {
					    System.out.println("  expr: " + statement + ", start pos: " + statement.getStartPosition() + ", get line num: "
							+ this.unit.getLineNumber(statement.getStartPosition()));
					}
					MethodInvocation methodInvoke = (MethodInvocation)expression;
					//init a statement
					MethodInvocation tracerInvoke = methodInvoke.getAST().newMethodInvocation();
					tracerInvoke.setExpression(methodInvoke.getAST().newName(TRACER_CLASS));
					tracerInvoke.setName(methodInvoke.getAST().newSimpleName(RECORD_OBSERVATION));
					List<Expression> arguments = tracerInvoke.arguments();
//					arguments.add(methodInvoke.getAST().newNumberLiteral( String.valueOf(this.unit.getLineNumber(methodInvoke.getStartPosition()))));
					arguments.add(methodInvoke.getAST().newNumberLiteral( String.valueOf(methodInvoke.getStartPosition())));
					StringLiteral stringLiteral = methodInvoke.getAST().newStringLiteral();
					stringLiteral.setLiteralValue(methodInvoke.arguments().get(0).toString());
					arguments.add(stringLiteral);
					Statement insertStatement = methodInvoke.getAST().newExpressionStatement(tracerInvoke);
					//add to the map
					insertingStatements.put(i, insertStatement);
					
					//add to the collection
					this.instrumentedObservers.add(methodInvoke);
				} else if (expression instanceof MethodInvocation && isObserverEquality((MethodInvocation)expression)) {
					//record whether two expression are equal or not
					if(SketchTestOptions.verbose) {
					    System.out.println("  equality: " + statement + ", start pos: " + statement.getStartPosition() + ", get line num: "
							+ this.unit.getLineNumber(statement.getStartPosition()));
					}
					MethodInvocation methodInvoke = (MethodInvocation)expression;
					//init a statement
					MethodInvocation tracerInvoke = methodInvoke.getAST().newMethodInvocation();
					tracerInvoke.setExpression(methodInvoke.getAST().newName(TRACER_CLASS));
					tracerInvoke.setName(methodInvoke.getAST().newSimpleName(RECORD_OBSERVATION));
					List<Expression> arguments = tracerInvoke.arguments();
//					arguments.add(methodInvoke.getAST().newNumberLiteral( String.valueOf(this.unit.getLineNumber(methodInvoke.getStartPosition()))));
					arguments.add(methodInvoke.getAST().newNumberLiteral( String.valueOf(methodInvoke.getStartPosition())));
					Checker.checkTrue(methodInvoke.arguments().size() == 2, "The size of observer equality must be 2, not: "
							+ methodInvoke.arguments().size());
					
					String observed_value = methodInvoke.arguments().get(0).toString() + Observer.SEP_EQUAL + methodInvoke.arguments().get(1).toString();
					StringLiteral stringLiteral1 = methodInvoke.getAST().newStringLiteral();
					stringLiteral1.setLiteralValue(observed_value);
					arguments.add(stringLiteral1);
					Statement insertStatement = methodInvoke.getAST().newExpressionStatement(tracerInvoke);
					//add to the map
					insertingStatements.put(i, insertStatement);
				}
			}
			i++;
		}
		
		if(!insertingStatements.isEmpty()) {
			Integer[] indices = (Integer[])insertingStatements.keySet().toArray(new Integer[0]);
			Arrays.sort(indices);
			int length = indices.length;
			for(int index = length - 1; index > -1; index--) {
				Integer lineIndex = indices[index];
				Statement s = insertingStatements.get(lineIndex);
				statements.add(lineIndex + 1/*try*/, s);
			}
		}
		
		return super.visit(block);
	}
	
	public static String[][] OBSERVER_METHODS = new String[][]{new String[]{"Observer", "observeEquality"}, new String[]{"Observer", "observeExpr"}};
	
	public static boolean isObserverEquality(MethodInvocation methodInvoke) {
		if(shouldSkip(methodInvoke)) {
			return false;
		}
		if(methodInvoke.getExpression() == null) {
			return false;
		}
		
		String invokeClass = methodInvoke.getExpression().toString().trim();
		String invokeMethod = methodInvoke.getName().toString().trim();
		
		return invokeClass.equals("Observer") && invokeMethod.equals("observeEquality");
	}
	
	public static boolean isObserverExpression(MethodInvocation methodInvoke) {
		if(shouldSkip(methodInvoke)) {
			return false;
		}
		if(methodInvoke.getExpression() == null) {
			return false;
		}
		
		String invokeClass = methodInvoke.getExpression().toString().trim();
		String invokeMethod = methodInvoke.getName().toString().trim();
		
		return invokeClass.equals("Observer") && invokeMethod.equals("observeExpr");
	}
	
	public static boolean isObserverMethod(MethodInvocation methodInvoke) {
		
		if(shouldSkip(methodInvoke)) {
			return false;
		}
		if(methodInvoke.getExpression() == null) {
			return false;
		}
		
		String invokeClass = methodInvoke.getExpression().toString().trim();
		String invokeMethod = methodInvoke.getName().toString().trim();
		
		boolean isObserver = false;
		for(String[] ob_method_pair : OBSERVER_METHODS) {
			Checker.checkTrue(ob_method_pair.length == 2, "The length of ob method should be 2, not: " + ob_method_pair.length);
			if(invokeClass.equals(ob_method_pair[0]) && invokeMethod.equals(ob_method_pair[1])) {
				isObserver = true;
				break;
			}
		}
		return isObserver;
//		
//		boolean isObserveEquals = methodInvoke.getExpression().toString().trim().equals("Observer")
//            && methodInvoke.getName().toString().trim().equals("observeEquals");
//		boolean isObserveExpr = methodInvoke.getExpression().toString().trim().equals("Observer")
//            && methodInvoke.getName().toString().trim().equals("observeExpr");
//		return isObserveEquals || isObserveExpr;
	}
	
	public static boolean isWrappedTryCatch(TryStatement tryStatement) {
		int catchClauseNum = tryStatement.catchClauses().size();
		if(catchClauseNum != 1) {
			return false;
		}
		CatchClause catchClause = (CatchClause) tryStatement.catchClauses().get(0);
		SingleVariableDeclaration exceptionDecl = catchClause.getException();
		
		if(exceptionDecl.getType().toString().equals(SketchedTryCatchBlock.EXCEPTION_TYPE)
				&& exceptionDecl.getName().toString().equals(SketchedTryCatchBlock.EXCEPTION_NAME)) {
			return true;
		}
		
		return false;
	}
	
	public static boolean shouldSkip(MethodInvocation methodInvoke) {
		String str = methodInvoke.toString();
		if(str.trim().startsWith(TRACER_CLASS)) {
			return true;
		}
		return false;
	}
}